package com.example.zklock;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.locks.Lock;

public class ZKLock implements ZKDistributeLock, Watcher {

    private ZooKeeper zk = null;
    private String rootLockName = "/locks";
    // zk connect counter
    private CountDownLatch zkConnectCountDouwnLatch = new CountDownLatch(1);
    // zk wait lock counter
    private CountDownLatch waitCountDownLatch;
    // current node name
    private String curName;
    // the node which current node is listening
    private String waitName;
    // competing resource
    private String lockName;
    private int zkTimeOut = 30000;

    public ZKLock(String connectUrl, String lockName) throws Exception {
        try {
            zk = new ZooKeeper(connectUrl, zkTimeOut, this);
            this.lockName = lockName;
            zkConnectCountDouwnLatch.await();
            //判断根节点是否存在，不存在创建
            Stat stat = zk.exists(rootLockName, false);
            if (stat == null) {
                zk.create(rootLockName, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }

        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    /**
     * Watcher接口实现类
     * @param watchedEvent
     */
    @Override
    public void process(WatchedEvent watchedEvent) {
        Event.EventType eventType = watchedEvent.getType();
        if (eventType == Event.EventType.NodeDeleted) {
            // if node is not exist,don't wait it
            if (this.waitCountDownLatch != null) {
                this.waitCountDownLatch.countDown();
            }
            // connection is success
            else if (eventType == Event.EventType.None)
                zkConnectCountDouwnLatch.countDown();
        }
    }

    @Override
    public void lock() {
        if (tryLock()) {
            System.out.println("#####成功获取锁######");
            return;
        } else {
            waitForLock();
        }
    }

    @Override
    public void unLock() {
        System.out.println(Thread.currentThread().getName() + "unlock");
        try {
            zk.delete(curName, -1);
            curName = null;
            zk.close();
        } catch (InterruptedException e) {

            e.printStackTrace();
        } catch (KeeperException e) {

            e.printStackTrace();
        }
    }
    private void waitForLock() {
        try {
            Stat stat = zk.exists(rootLockName + "/" + waitName, true);
            if (stat != null) {
                this.waitCountDownLatch = new CountDownLatch(1);
                this.waitCountDownLatch.await();
                this.waitCountDownLatch = null;
                System.out.println(Thread.currentThread().getName() + "等到了lock");
            }
        } catch (KeeperException e) {

            e.printStackTrace();
        } catch (InterruptedException e) {

            e.printStackTrace();
        }

    }
    public boolean tryLock() {
        try {
            //创建锁节点  临时序列节点
            curName = zk.create(rootLockName + "/" + lockName + "_lock_", new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE,
                    CreateMode.EPHEMERAL_SEQUENTIAL);
            //获取根节点下的所有节点
            List<String> subNodes = zk.getChildren(rootLockName, false);
            Collections.sort(subNodes);
            // get all the children nodes compared with current node to judge it if is the smallest
            //最小节点代表是第一个，说明获得锁了
            if (curName.equals(rootLockName + "/" + subNodes.get(0))) {
                return true;
            }
            int curIndex = subNodes.indexOf(curName.split("/")[2]);
            // if the current node is not the smallest,find the one in front
            waitName = subNodes.get(curIndex - 1);
        } catch (KeeperException e) {

            e.printStackTrace();
        } catch (InterruptedException e) {

            e.printStackTrace();
        }
        return false;
    }


}
